"""
FileName: dcm2nii.py
Function: 将目录下的所有子目录中的dcm分别打包成nii文件
"""
import os
import pydicom
import numpy as np
import nibabel as nib
import SimpleITK as sitk


# 获取dcm文件的序列信息
def get_series_info_sitk(dicom_file_path):
    # 读取DICOM文件
    dicom_reader = sitk.ImageFileReader()
    dicom_reader.SetFileName(dicom_file_path)
    dicom_reader.ReadImageInformation()

    # 获取Series Description和Series Number
    series_description = dicom_reader.GetMetaData("0008|103e")  # Series Description
    series_number = dicom_reader.GetMetaData("0020|0011")  # Series Number

    return series_description.strip(), series_number


def get_series_info_py(dicom_file_path):
    try:
        # 读取DICOM文件
        dicom_dataset = pydicom.dcmread(dicom_file_path)

        # 获取Series Number和Series Description
        series_number = dicom_dataset.SeriesNumber
        series_description = dicom_dataset.SeriesDescription

        return series_description, series_number
    except Exception as e:
        print(f"Error: {e}")
        return None, None


# 获取input_folder文件夹下的dcm图片序列
def get_nii_name(dicom_file_path):
    dicom_files = [f for f in os.listdir(dicom_file_path) if f.endswith('.dcm')]
    if not dicom_files:
        print("文件夹中没有DICOM文件")
        return

    # 选择文件夹中的第一个DICOM文件, 获取Series Description和Series Number
    first_dcm_file = os.path.join(dicom_file_path, dicom_files[0])
    series_description, series_number = get_series_info_py(first_dcm_file)

    # 构建新的文件名 "SeriesNumber-SeriesDescription.nii"
    # series_number = str(series_number).zfill(2)  # 序列号左侧填0
    nii_name = f"{series_number}-{series_description}.nii"
    return nii_name


def dicom_series_to_nifti(dcm_path, nii_path, standard_size=0):
    """将dcm_path文件夹下的多个dcm图片序列合成为1个nii文件，存放在nii_path路径下

    :param dcm_path: dcm源文件路径
    :param nii_path: nii文件存放路径
    :param standard_size: 目标 nii 文件边长大小
    """
    dicom_files = [f for f in os.listdir(dcm_path) if f.endswith('.dcm')]
    if not dicom_files:  # 检查输入路径是否有dcm文件
        print("文件夹中没有DICOM文件")
        return
    if not os.path.exists(nii_path):  # 检查输出路径是否存在
        os.makedirs(nii_path)

    first_dcm_file = os.path.join(dcm_path, dicom_files[0])  # 取dcm序列第1个dcm文件作为信息源
    dicom_file_info = pydicom.dcmread(first_dcm_file)
    series_number = dicom_file_info.SeriesNumber  # 获取序列号和序列描述
    series_description = dicom_file_info.SeriesDescription
    study_id = dicom_file_info.StudyID
    nii_filename = f"{series_number}-{series_description}.nii"
    temp_file_path = os.path.join(nii_path, "temp.nii")

    dicom_reader = sitk.ImageSeriesReader()
    dicom_series = dicom_reader.GetGDCMSeriesFileNames(dcm_path)  # 获取文件夹中dicom文件列表
    dicom_reader.SetFileNames(dicom_series)
    image = dicom_reader.Execute()  # 执行读取

    original_size = image.GetSize()
    print("before_size:", original_size)
    need_standard_size = standard_size if standard_size != 0 else original_size[0]  # 获取设定的标准值
    pad_height = (need_standard_size - original_size[1]) // 2  # 计算填充量
    pad_weight = (need_standard_size - original_size[0]) // 2
    global padded_image_array
    try:
        padded_image_array = np.pad(sitk.GetArrayFromImage(image), (
            (0, 0), (pad_height, pad_height), (pad_weight, pad_weight)
        ), mode='constant', constant_values=0)  # 使用np.pad对图像进行上下填充
    except Exception as e:
        print('padded Error:', e)
    image = sitk.GetImageFromArray(padded_image_array)
    print("after_size :", image.GetSize())

    nifti_image = sitk.Cast(image, sitk.sitkFloat32)  # 将读取的image转化为NIfTI图像
    sitk.WriteImage(nifti_image, temp_file_path)  # 保存临时 NIfTI 文件

    nifti_img = nib.load(temp_file_path)
    nifti_img.header['descrip'] = f"Study ID: {study_id}"  # 保存nii文件的studyID信息
    nii_file_path = os.path.join(nii_path, nii_filename)
    nib.save(nifti_img, nii_file_path)
    print(nii_file_path)
    if os.path.exists(temp_file_path):
        os.remove(temp_file_path)  # 删除临时文件



# 获取 dcm2nii时封装的Study ID
def get_nii_description(nii_file_path):
    global header_info
    try:
        nii_img = nib.load(nii_file_path)
        header_info = nii_img.header
        if header_info:
            import numpy as np
            patient_id = np.array2string(header_info.get('descrip', ''))
            study_id = (patient_id.split('Study ID: ')[-1]).split("'")[0]
            return study_id
        else:
            print("Failed to get header information.")
    except Exception as e:
        print("Error:", e)


if __name__ == "__main__":
    # 以处理2022HCC数据集的11个患者序列为例
    all_path = r"C:\Users\18049\Desktop\nii"
    dcm_path_2022 = all_path + r"\1_unclassified_dcm\metastasis 2022"
    nii_path_2022 = all_path + r"\2_classification_nii\metastasis 2022"
    # 遍历根文件夹下的所有子文件夹
    # 目标结构：metastasis 2022\XXX_M000xxxxx\MR__20XX0X0Xxxxxxx_1XXXX2\15_t1_vibe_fs_tra_bh_p2_two scans_2XXX8
    for patient_dcm_folder in os.listdir(dcm_path_2022):
        patient_dcm_folder_path = os.path.join(dcm_path_2022, patient_dcm_folder)  # C:\..\CHEN YI JIANG_M00xxxxx
        if os.path.isdir(patient_dcm_folder_path):  # 检查是否是文件夹
            print('-' * 50, '\n', patient_dcm_folder, sep='')
            patient_nii_folder_path = os.path.join(nii_path_2022, patient_dcm_folder)  # 创建存放目标nii的文件夹
            if not os.path.exists(patient_nii_folder_path):
                os.makedirs(patient_nii_folder_path)
            mr_temp_folder = os.listdir(patient_dcm_folder_path)[0]
            mr_temp_folder_path = os.path.join(patient_dcm_folder_path, mr_temp_folder)  # C:\..\MR__20220xxxx55_19xxxxx
            for phases_folder in os.listdir(mr_temp_folder_path):
                phases_folder_path = os.path.join(mr_temp_folder_path, phases_folder)
                # print(phases_folder_path)
                try:
                    dicom_series_to_nifti(phases_folder_path, patient_nii_folder_path)
                except Exception as e:
                    # 处理异常，例如输出错误信息
                    print(f"Error processing {phases_folder}: {str(e)}")
                    # 继续循环到下一个对象
                    continue

    # --------------------------------------------------------------------------------------------------------------
    # 单次dcm转换nii测试
    # input = r"C:\Users\18049\Desktop\nii\1_unclassified_dcm\metastasis 2022" \
    #         r"\CHEN YI JIANG_M00043899\MR__20220304112655_1931812\15_t1_vibe_fs_tra_bh_p2_two scans_23188"
    # output = r"C:\Users\18049\Desktop\nii\2_classification_nii\metastasis 2022"
    # dicom_series_to_nifti(input, output)
    #
    # file_path = r"C:\Users\18049\Desktop\nii\2_classification_nii\metastasis 2022\15-t1_vibe_fs_tra_bh_p2_two scans.nii"
    # print("Study ID: ", get_nii_description(file_path))